""" Options
"""

import argparse
import os
import torch


class Options():
    """Options class

    Returns:
        [argparse]: argparse containing train and test options
    """

    def __init__(self):
        # Inputs for the main function
        self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

        # 不修改的数据
        self.parser.add_argument(
            '--data_name',
            choices=['sine', 'stock', 'energy', 'water'],
            default='water',
            type=str)
        self.parser.add_argument(
            '--masked_ratio',
            help='mask ratio',
            default=0.3,  # 如果为空，读取masked_water_data.csv，所有指标都被mask而不是某个指标
            type=float)
        self.parser.add_argument(  # 这个没有用了，这个是根据输入数据来决定的,但是需要保留这个参数用于模型的初始化
            '--z_dim',
            help='z or data dimension',
            default=6,
            type=int)
        self.parser.add_argument(
            '--module',
            choices=['gru', 'lstm', 'lstmLN'],
            help='only gru, grui using other param',
            default='gru',
            type=str)
        self.parser.add_argument(
            '--log_interval',
            help='log per iter',
            default=100,
            type=int)
        self.parser.add_argument(
            '--metric_iteration',
            default=10,
            type=int
        )
        # 需要修改的数据参数
        self.parser.add_argument(
            '--masked_indicator',
            help='mask indicator if None mask all indicator',
            default='',  # 如果为空，读取masked_water_data.csv，所有指标都被mask而不是某个指标
            type=str)
        self.parser.add_argument(
            '--indicators',
            help='input water indicators',
            default='WATER_TEMPERATURE,PH_VALUE,TOTAL_NITROGEN,DISSOLVED_OXYGEN',
            type=str)
        self.parser.add_argument(
            '--data_type',
            type=str,
            choices=['masked', 'ori'],
            help='train data type, masked or ori',
            default='masked')
        self.parser.add_argument(
            '--delta_type',
            type=str,
            choices=['normal', 'time_delta', 'zero'], # normal对输入0，不做特殊处理，time_delta输入0，则将上一个有效值看作上一个值，
            help='',
            default='time_delta')

        # 需要修改的模型参数
        self.parser.add_argument(
            '--gru_type',
            choices=['gru', 'grui'],
            default='grui',
            type=str)
        self.parser.add_argument(
            '--seq_len',
            help='sequence length',
            default=12,
            type=int)
        self.parser.add_argument(
            '--hidden_dim',
            help='hidden state dimensions (should be optimized)',
            default=64,
            type=int)
        self.parser.add_argument(
            '--num_layer',
            help='number of layers (should be optimized)',
            default=2,
            type=int)
        self.parser.add_argument(
            '--dropout',
            type=float,
            help='dropout',
            default=0.1)
        self.parser.add_argument(
            '--layer_size',
            choices=['hidden_size', 'z_dim'],
            default='z_dim',
            type=str)
        self.parser.add_argument(
            '--is_ln',
            default='True',
            type=str)
        self.parser.add_argument(
            '--is_fc',
            default='True',
            help='if z_dim must True',  # 如果是layer_size=z_dim 必须True。因为要将hidden_size转化为output_size=z_dim
            type=str)
        self.parser.add_argument(
            '--is_res',
            default='True',  # 不必须有fc，但是只有input=output_size fc才会起作用
            type=str)
        self.parser.add_argument(
            '--g_loss_weight',
            help='调整 （生成和原始mean，std差距） 与 （d输出的生成数据的真实程度）之间的权重',
            default=100,
            type=int)
        self.parser.add_argument(
            '--n_head',
            help='多头注意力的头',
            default=8,
            type=int)
        self.parser.add_argument(
            '--snr',
            help='数据增强',
            default=50,
            type=int)

        # 需要修改的训练参数
        self.parser.add_argument(
            '--is_load',
            help='是否加载已有模型，如果加载失败再训练',
            default='True',
            type=str)
        self.parser.add_argument(
            '--is_parallel',
            help='是否加载已有模型，如果加载失败再训练',
            default='False',
            type=str)
        self.parser.add_argument(
            '--iteration',
            help='Training iterations (should be optimized)',
            default=5,  # 50000
            type=int)
        self.parser.add_argument(
            '--impute_iter',
            help='Imputing iterations (should be optimized)',
            default=1,  # 5000
            type=int)
        self.parser.add_argument(
            '--batch_size',
            help='the number of samples in mini-batch (should be optimized)',
            default=128,
            type=int)

        # 需要修改的其他数据
        self.parser.add_argument(
            '--output_dir',
            default='output_v6',
            help='folder to output images and model checkpoints')
        self.parser.add_argument(
            '--run_id',
            default='',
            help='')
        self.parser.add_argument(
            '--manualseed',
            default=2020,
            type=int,
            help='manual seed')
        self.parser.add_argument(
            '--resume',
            default='',
            help="path to checkpoints (to continue training)")


        self.parser.add_argument('--GAN_type', type=str, help='gan type', default='gan')

        # 一些默认的参数，
        self.parser.add_argument('--is_normalize', type=bool, help='normalize or not', default=True)
        self.parser.add_argument('--workers', type=int, help='number of data loading workers', default=8)
        self.parser.add_argument('--device', type=str, default='gpu', help='Device: gpu | cpu')
        self.parser.add_argument('--gpu_ids', type=str, default='0,1,2', help='可用的GPU: e.g. 0  0,1,2, 0,2. use -1 for CPU')
        self.parser.add_argument('--gpu_id', type=str, default='0', help='number of GPUs to use')
        self.parser.add_argument('--model', type=str, default='TimeGAN', help='chooses which model to use. timegan')
        self.parser.add_argument('--mode', type=str, default='train', choices=['train', 'test', 'impute'], help='run mode, options[train, test, impute]')
        self.parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment')
        self.parser.add_argument('--display_server', type=str, default="http://localhost", help='visdom server of the web display')
        self.parser.add_argument('--display_port', type=int, default=8097, help='visdom port of the web display')
        self.parser.add_argument('--display_id', type=int, default=0, help='window id of the web display')
        self.parser.add_argument('--display', action='store_true', help='Use visdom.')

        # 不修改的模型训练参数
        self.parser.add_argument('--beta1', type=float, default=0.9, help='momentum term of adam')
        self.parser.add_argument('--lr', type=float, default=0.001, help='initial learning rate for adam')

        self.parser.add_argument('--w_gamma', type=float, default=1, help='Gamma weight') # D的输入的假样本AR和生成权重
        self.parser.add_argument('--w_es', type=float, default=0.1, help='Encoder loss weight')
        self.parser.add_argument('--w_e0', type=float, default=10, help='Encoder loss weight')
        self.parser.add_argument('--w_g', type=float, default=100, help='Generator loss weight.')
        # 超参
        # 论文中的gamma # todo 这个参数有用吗，插值的时候d_loss很大，但是论文中说对于预测任务lamnda为0效果最优，分类任务不说
        self.parser.add_argument('--w_lambda', type=float, default=0.0, help='discriminator loss weight in impute loss')
        self.parser.add_argument('--w_eta', type=float, default=100, help='L_AR and L_G weight L_AR weight')
        self.parser.add_argument('--w_zeta', type=float, default=100, help='L_G and reconstruction weight, reconstruction weight')
        self.parser.add_argument('--w_lambda2', type=float, default=0.01, help='L_AR and L_R, L_AR weight')

        # WGAN参数
        self.parser.add_argument('--w_clip_c', type=float, default=0.01, help='wgan clip c')
        self.parser.add_argument('--w_lr', type=float, default=0.00005, help='rmsprop learning rate')
        self.parser.add_argument('--w_d_iter', type=int, default=5, help='per adjust G how times adjust D')

        self.parser.add_argument('--dis_iteration', type=int, default=5000, help='dis_score iteration')

        self.isTrain = True
        self.opt = None

    def parse(self):
        """ Parse Arguments.
        """
        self.opt = self.parser.parse_args()
        self.opt.isTrain = self.isTrain  # train or test
        # set gpu ids
        if self.opt.device == 'gpu':
            os.environ["CUDA_VISIBLE_DEVICES"] = self.opt.gpu_ids
            torch.cuda.set_device(int(self.opt.gpu_id))

        return self.opt